#当我想检测绳子的时候，我发现yolo的方法可能不太行，所以想尝试一下其他方法进行增强
#差分是其中一种
#该脚本是图片差分的增强版，用于对目录下所有图片进行差分

import argparse
import cv2
import numpy as np
import os
import json

def Parse_Arguments():
    parser = argparse.ArgumentParser(description="图片差分")
    parser.add_argument("--background-dir", type=str, default=r"D:\tempDownload\dataset\elevator\rope_glass_rope_small_voc\background")
    parser.add_argument("--source-dir", type=str, default=r"D:\tempDownload\dataset\elevator\rope_glass_voc\JPEGImages")
    parser.add_argument("--target-dir", type=str, default=r"D:\tempDownload\dataset\elevator\rope_glass_voc\Residual")
    return parser.parse_args()


def main():
    args = Parse_Arguments()
    background_dir = args.background_dir
    background_filenames = os.listdir(background_dir)
    source_dir = args.source_dir
    target_dir = args.target_dir
    source_filenames = os.listdir(source_dir)
    background_map = {}
    with open(os.path.join(background_dir, "map.json"), "r", encoding="utf-8") as json_file:
            background_map = json.load(json_file)

    for background_filename in background_filenames:
        if not background_filename.lower().endswith((".png", ".jpg")):
            break
        background_img = cv2.imread(os.path.join(background_dir, background_filename))
        background_prefix = background_filename.split('.')[0]
        source_prefixs = background_map[background_prefix]
        for source_prefix in source_prefixs:
            for source_filename in list(source_filenames):
                source_gt_prefix = int(source_filename.split('.')[0].split('_')[0])
                if int(source_prefix)==source_gt_prefix:
                    source_img = cv2.imread(os.path.join(source_dir, source_filename))
                    target_height, target_width = source_img.shape[:2]
                    background_height, background_width = background_img.shape[:2]
                    if target_height != background_height or target_width != background_width:
                        background_img_resize = cv2.resize(background_img, (target_width, target_height))
                        print("发生resize")
                    else:
                        background_img_resize = background_img
                    diff_image = cv2.absdiff(background_img_resize, source_img)
                    diff_image = cv2.cvtColor(diff_image, cv2.COLOR_BGR2GRAY)
                    # print(diff_image.shape)
                    target_filename = source_filename.split('.')[0] + '.jpg'
                    source_filenames.remove(source_filename)
                    cv2.imwrite(os.path.join(target_dir, target_filename), diff_image)
                    print(f"{os.path.join(target_dir, target_filename)} done!!!")
                    print(f"对应的背景图为{background_filename}")
                    print(f"还剩{len(source_filenames)}")
    last_remove_filename = ""
    for source_filename in list(source_filenames):
        source_img = cv2.imread(os.path.join(source_dir, source_filename))
        edges = cv2.Canny(source_img, 50, 150)
        target_filename = source_filename.split('.')[0] + '.jpg'
        cv2.imwrite(os.path.join(target_dir, target_filename), edges)
        # if last_remove_filename != "":
        #     source_filenames.remove(last_remove_filename)
        #     last_remove_filename = source_filename
        # else:
        #     last_remove_filename = source_filename
        source_filenames.remove(source_filename)
        print(f"{os.path.join(target_dir, target_filename)} done!!!")
        print(f"还剩{len(source_filenames)}")

if __name__ == "__main__":
    main()
